package com.feriki.msync.utils;

import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DescribeCollectionResponse;
import io.milvus.grpc.DescribeIndexResponse;
import io.milvus.grpc.IndexDescription;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.DescribeCollectionParam;
import io.milvus.param.collection.DropCollectionParam;
import io.milvus.param.collection.HasCollectionParam;
import io.milvus.param.collection.LoadCollectionParam;
import io.milvus.param.index.DescribeIndexParam;
import io.milvus.response.DescCollResponseWrapper;

import java.util.List;

/**
 * @author luozhida
 * @create 2024-07-01
 */
public class MilvusOperatorUtil {

    public static boolean checkCollectionExist(String collectionName , MilvusServiceClient milvusClient){
        R<Boolean> booleanR = milvusClient.hasCollection(HasCollectionParam.newBuilder().withCollectionName(collectionName).build());
        return booleanR.getData() == Boolean.TRUE ? Boolean.TRUE : Boolean.FALSE;
    }

    public static DescCollResponseWrapper getCollectionDescription(String collectionName , MilvusServiceClient milvusClient){
        R<DescribeCollectionResponse> responseR = milvusClient.describeCollection(DescribeCollectionParam.newBuilder().withCollectionName(collectionName).build());
        if(R.Status.Success.getCode() == responseR.getStatus() && responseR.getData()!=null){
            return new DescCollResponseWrapper(responseR.getData());
        }
        return null;
    }

    public static boolean dropCollection(String collectionName , MilvusServiceClient milvusClient){
        R<RpcStatus> statusR = milvusClient.dropCollection(DropCollectionParam.newBuilder().withCollectionName(collectionName).build());
        return R.Status.Success.getCode() == statusR.getStatus();
    }

    public static boolean loadCollection(String collectionName , MilvusServiceClient milvusClient){
        R<RpcStatus> collectionLoadStatus = milvusClient.loadCollection(
                LoadCollectionParam.newBuilder()
                        .withCollectionName(collectionName)
                        .build()
        );
        return R.Status.Success.getCode() == collectionLoadStatus.getStatus();
    }

    public static List<IndexDescription> getIndexes(String collectionName , MilvusServiceClient milvusClient){
        R<DescribeIndexResponse> describeIndexResponseR = milvusClient.describeIndex(DescribeIndexParam.newBuilder().withCollectionName(collectionName).build());
        if(R.Status.Success.getCode() == describeIndexResponseR.getStatus() && describeIndexResponseR.getData()!=null){
            return describeIndexResponseR.getData().getIndexDescriptionsList();
        }
        return null;
    }

}
